# ai_premium_predictor/data/stock_data_loader.py

from sqlalchemy import create_engine
import pandas as pd
from config import CONFIG

class StockDataLoader:
    def __init__(self):
        db_cfg = CONFIG["db"]
        # 使用 SQLAlchemy 创建连接引擎
        self.engine = create_engine(
            f"mysql+pymysql://{db_cfg['user']}:{db_cfg['password']}@{db_cfg['host']}:{db_cfg['port']}/{db_cfg['database']}?charset=utf8mb4"
        )

    def get_deal_dates(self):
        sql = "SELECT add_date FROM deal_date ORDER BY add_date"
        result = pd.read_sql(sql, self.engine)  # 使用 SQLAlchemy 引擎
        print(f"[LOG] Query executed: {sql}, Rows fetched: {result.shape[0]}")
        return result

    def get_hit_board_data(self, start_date=None, end_date=None):
        sql = "SELECT * FROM stock_hit_board WHERE 1=1"
        if start_date:
            sql += f" AND add_date >= '{start_date}'"
        if end_date:
            sql += f" AND add_date <= '{end_date}'"
        result = pd.read_sql(sql, self.engine)  # 使用 SQLAlchemy 引擎
        print(f"[LOG] Query executed: {sql}, Rows fetched: {result.shape[0]}")
        return result

    def get_daily_summary_data(self, start_date=None, end_date=None):
        sql = "SELECT * FROM daily_stock_summary WHERE 1=1"
        if start_date:
            sql += f" AND summary_date >= '{start_date}'"
        if end_date:
            sql += f" AND summary_date <= '{end_date}'"
        result = pd.read_sql(sql, self.engine)  # 使用 SQLAlchemy 引擎
        print(f"[LOG] Query executed: {sql}, Rows fetched: {result.shape[0]}")
        return result